# src/visualization/plots.py
import plotly.graph_objects as go
import plotly.express as px
from plotly.subplots import make_subplots
import pandas as pd
import numpy as np
from typing import List, Dict, Any
import json

class ParetoFrontVisualizer:
    """Pareto前沿可视化器"""
    
    def __init__(self):
        self.colors = ['#1f77b4', '#ff7f0e', '#2ca02c', '#d62728', '#9467bd']
    
    def create_pareto_visualization(self, pareto_set: List[Dict[str, Any]], 
                                  user_preferences: Dict[str, float]) -> Dict[str, Any]:
        """创建Pareto前沿可视化"""
        
        # 提取数据
        solutions_data = []
        for i, solution in enumerate(pareto_set):
            solutions_data.append({
                'id': i,
                'return': solution['metrics']['expected_return'],
                'risk': solution['metrics']['risk'],
                'esg': solution['metrics'].get('esg_score', 0),
                'weights': solution['weights']
            })
        
        df = pd.DataFrame(solutions_data)
        
        # 创建3D散点图
        scatter_3d = self._create_3d_scatter(df)
        
        # 创建平行坐标图
        parallel_coords = self._create_parallel_coordinates(df)
        
        # 创建权重分布图
        weights_chart = self._create_weights_distribution(df)
        
        # 创建效率前沿图
        efficient_frontier = self._create_efficient_frontier(df)
        
        return {
            'scatter_3d': scatter_3d,
            'parallel_coordinates': parallel_coords,
            'weights_distribution': weights_chart,
            'efficient_frontier': efficient_frontier,
            'summary_stats': self._calculate_summary_stats(df)
        }
    
    def _create_3d_scatter(self, df: pd.DataFrame) -> str:
        """创建3D散点图"""
        fig = go.Figure(data=[go.Scatter3d(
            x=df['return'],
            y=df['risk'],
            z=df['esg'],
            mode='markers',
            marker=dict(
                size=8,
                color=df['return'],
                colorscale='Viridis',
                showscale=True,
                colorbar=dict(title="期望收益")
            ),
            text=[f"解{i}: 收益={r:.3f}, 风险={risk:.3f}, ESG={esg:.1f}" 
                  for i, (r, risk, esg) in enumerate(zip(df['return'], df['risk'], df['esg']))],
            hovertemplate='<b>解 %{text}</b><br>' +
                         '期望收益: %{x:.3f}<br>' +
                         '风险: %{y:.3f}<br>' +
                         'ESG评分: %{z:.1f}<br>' +
                         '<extra></extra>'
        )])
        
        fig.update_layout(
            title='📊 Pareto前沿 3D可视化',
            scene=dict(
                xaxis_title='期望收益',
                yaxis_title='风险',
                zaxis_title='ESG评分'
            ),
            width=800,
            height=600
        )
        
        return fig.to_json()
    
    def _create_parallel_coordinates(self, df: pd.DataFrame) -> str:
        """创建平行坐标图"""
        fig = go.Figure(data=go.Parcoords(
            line=dict(color=df['return'], colorscale='Viridis', showscale=True),
            dimensions=list([
                dict(range=[df['return'].min(), df['return'].max()],
                     label='期望收益', values=df['return']),
                dict(range=[df['risk'].min(), df['risk'].max()],
                     label='风险', values=df['risk']),
                dict(range=[df['esg'].min(), df['esg'].max()],
                     label='ESG评分', values=df['esg'])
            ])
        ))
        
        fig.update_layout(
            title='📈 目标函数平行坐标图',
            width=800,
            height=400
        )
        
        return fig.to_json()
    
    def _create_weights_distribution(self, df: pd.DataFrame) -> str:
        """创建权重分布图"""
        # 获取所有资产名称
        all_assets = set()
        for weights in df['weights']:
            all_assets.update(weights.keys())
        
        # 创建权重矩阵
        weight_matrix = []
        for weights in df['weights']:
            weight_row = [weights.get(asset, 0) for asset in sorted(all_assets)]
            weight_matrix.append(weight_row)
        
        weight_df = pd.DataFrame(weight_matrix, columns=sorted(all_assets))
        
        # 创建堆叠条形图
        fig = go.Figure()
        
        for i, asset in enumerate(sorted(all_assets)):
            fig.add_trace(go.Bar(
                name=asset,
                x=list(range(len(weight_df))),
                y=weight_df[asset],
                marker_color=self.colors[i % len(self.colors)]
            ))
        
        fig.update_layout(
            title='💼 投资组合权重分布',
            xaxis_title='解的序号',
            yaxis_title='权重',
            barmode='stack',
            width=800,
            height=500
        )
        
        return fig.to_json()
    
    def _create_efficient_frontier(self, df: pd.DataFrame) -> str:
        """创建效率前沿图"""
        fig = go.Figure()
        
        # 绘制效率前沿
        fig.add_trace(go.Scatter(
            x=df['risk'],
            y=df['return'],
            mode='markers+lines',
            name='效率前沿',
            marker=dict(
                size=10,
                color=df['esg'],
                colorscale='RdYlGn',
                showscale=True,
                colorbar=dict(title="ESG评分")
            ),
            line=dict(color='blue', width=2),
            text=[f"ESG: {esg:.1f}" for esg in df['esg']],
            hovertemplate='风险: %{x:.3f}<br>' +
                         '收益: %{y:.3f}<br>' +
                         '%{text}<br>' +
                         '<extra></extra>'
        ))
        
        fig.update_layout(
            title='📈 风险-收益效率前沿',
            xaxis_title='风险',
            yaxis_title='期望收益',
            width=800,
            height=500,
            showlegend=True
        )
        
        return fig.to_json()
    
    def _calculate_summary_stats(self, df: pd.DataFrame) -> Dict[str, Any]:
        """计算汇总统计"""
        return {
            'total_solutions': len(df),
            'return_range': {
                'min': float(df['return'].min()),
                'max': float(df['return'].max()),
                'mean': float(df['return'].mean())
            },
            'risk_range': {
                'min': float(df['risk'].min()),
                'max': float(df['risk'].max()),
                'mean': float(df['risk'].mean())
            },
            'esg_range': {
                'min': float(df['esg'].min()),
                'max': float(df['esg'].max()),
                'mean': float(df['esg'].mean())
            },
            'best_sharpe_solution': int(df.loc[df['return'] / df['risk'] == (df['return'] / df['risk']).max()].index[0]),
            'best_esg_solution': int(df.loc[df['esg'] == df['esg'].max()].index[0])
        }

class DashboardCreator:
    """仪表盘创建器"""
    
    def __init__(self):
        self.visualizer = ParetoFrontVisualizer()
    
    def create_optimization_dashboard(self, optimization_results: Dict[str, Any]) -> str:
        """创建优化结果仪表盘"""
        
        # 创建子图
        fig = make_subplots(
            rows=2, cols=2,
            subplot_titles=('实时监控', '风险指标', '收益分析', 'ESG评分'),
            specs=[[{"secondary_y": True}, {"type": "scatter"}],
                   [{"type": "bar"}, {"type": "pie"}]]
        )
        
        # 实时监控
        dates = pd.date_range(start='2023-01-01', periods=100, freq='D')
        portfolio_values = np.cumsum(np.random.normal(0.001, 0.02, 100)) + 1
        
        fig.add_trace(
            go.Scatter(x=dates, y=portfolio_values, name='组合净值'),
            row=1, col=1
        )
        
        # 风险指标
        risk_metrics = ['VaR', 'CVaR', '最大回撤', '波动率']
        risk_values = [0.03, 0.05, 0.08, 0.15]
        
        fig.add_trace(
            go.Bar(x=risk_metrics, y=risk_values, name='风险指标'),
            row=1, col=2
        )
        
        # 收益分析
        periods = ['1个月', '3个月', '6个月', '1年']
        returns = [0.02, 0.08, 0.15, 0.25]
        
        fig.add_trace(
            go.Bar(x=periods, y=returns, name='历史收益'),
            row=2, col=1
        )
        
        # ESG评分分布
        sectors = ['科技', '金融', '医疗', '消费', '能源']
        esg_scores = [85, 72, 90, 68, 45]
        
        fig.add_trace(
            go.Pie(labels=sectors, values=esg_scores, name='ESG评分'),
            row=2, col=2
        )
        
        fig.update_layout(
            title='🎯 投资组合优化仪表盘',
            height=800,
            showlegend=True
        )
        
        return fig.to_json()
